import matplotlib.pyplot as plt
import numpy as np
from tqdm import tqdm
from sarsa import CliffWalkingEnv, print_agent  
import random
import time

class DynaQ:
    def __init__(self, ncol, nrow, epsilon, alpha, gamma, n_planning, n_action=4):
        self.Q = np.zeros([nrow * ncol, n_action])
        self.n_action = n_action
        self.alpha = alpha
        self.gamma = gamma
        self.epsilon = epsilon

        self.n_planning = n_planning
        self.model = dict()

    def take_action(self, state):
        if np.random.random() < self.epsilon:
            action = np.random.randint(self.n_action)   
        else:
            action = np.argmax(self.Q[state])
        return action

    def q_update(self, s0, a0, r, s1):
        td_error = r + self.gamma * self.Q[s1].max() - self.Q[s0, a0]
        self.Q[s0, a0] += self.alpha * td_error

    def update(self, s0, a0, r, s1):
        self.q_update(s0, a0, r, s1)
        self.model[(s0, a0)] = r, s1
        for _ in range(self.n_planning):
            # 随机选择曾经遇到过的状态动作对
            (s, a), (r, s_) = random.choice(list(self.model.items()))
            self.q_update(s, a, r, s_)

def DynQ_episodes(agent, env, num_episodes, show=True):
    return_list = []
    for i in range(10):
        with tqdm(total=int(num_episodes/10), desc='Iteration %d' % i) as pbar:
            for i_episode in range(int(num_episodes/10)):
                episode_return = 0
                state = env.reset()
                done = False
                while not done:
                    action = agent.take_action(state)
                    next_state, reward, done = env.step(action)
                    episode_return += reward
                    agent.update(state, action, reward, next_state)
                    state = next_state

                return_list.append(episode_return)
                if (i_episode+1) % 10 == 0:
                    pbar.set_postfix({
                        'episode':
                        '%d' % (num_episodes / 10 * i + i_episode + 1),
                        'return':
                        '%.3f' % np.mean(return_list[-10:])
                    })
                pbar.update(1)

    if show:
        episodes_list = list(range(len(return_list)))
        plt.plot(episodes_list, return_list, label = str(agent.n_planning) + ' planning steps')
        

np.random.seed(0)
random.seed(0)

n_planning_list = [0, 2, 20]

ncol = 12
nrow = 4
env = CliffWalkingEnv(ncol, nrow)
alpha = 0.1
epsilon = 0.01
gamma = 0.9
num_episodes = 500
for n_planning in n_planning_list:
    print("n_planning steps = %d" % n_planning)
    agent = DynaQ(ncol, nrow, epsilon, alpha, gamma, n_planning)
    DynQ_episodes(agent, env, num_episodes)
    time.sleep(0.5)

plt.legend()
plt.xlabel('Episodes')
plt.ylabel('Returns')
plt.title('Dyna-Q on {}'.format('Cliff Walking'))
plt.show()